from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.spatial import KDTree
from sklearn.linear_model import LinearRegression
import umap
from ripser import ripser
from persim import plot_diagrams
from scipy.spatial.distance import pdist, squareform


def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal

def main():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration

    nn_type = ''
    if rpl_config.nn_type == "vanilla":
        nn_type = "vanilla"
    elif rpl_config.nn_type == "gru":
        nn_type = "gru"

    """ 载入混沌轨迹
    """
    file_name = "./logs/chaoitc_trajectories_" + rpl_config.nn_type + "_" + str(rpl_config.nn_size) + ".npy"
    
    chaoitc_trajectories = np.load(file_name)

    print(chaoitc_trajectories.shape)

    # 随机选择一个 trajectory
    rnd_idx = np.random.randint(chaoitc_trajectories.shape[0])
    print("rnd_idx: ", rnd_idx)

    trj_chosen = chaoitc_trajectories[rnd_idx]

    # 将 trj_chosen 进行 PCA，然后绘制到一个三维图中
    pca1 = PCA()
    pca1.fit(trj_chosen)

    # 打印 variance ratio
    print("pca1.explained_variance_ratio_: ", pca1.explained_variance_ratio_)

    trj_chosen_pca = pca1.transform(trj_chosen)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    # ax.scatter(trj_chosen_pca[:, 0], trj_chosen_pca[:, 1], trj_chosen_pca[:, 2], s=1)
    # 绘制连线图，并且指定透明度为0.1
    ax.plot(trj_chosen_pca[:, 0], trj_chosen_pca[:, 1], trj_chosen_pca[:, 2], alpha=1)
    # for i in range(trj_chosen_pca.shape[0] - 1):
    #     ax.plot(trj_chosen_pca[i:i+2, 0], trj_chosen_pca[i:i+2, 1], trj_chosen_pca[i:i+2, 2],c='blue', alpha=0.1)

    plt.show()

            
if __name__ == "__main__":
    main()
